nanopyx.core.analysis.rcc

  1# REF: based on https://github.com/jungmannlab/picasso/blob/d867f561ffeafce752f37968a20698556d04dafb/picasso/imageprocess.py
  2
  3import numpy as np
  4import lmfit
  5from tqdm import tqdm
  6
  7from .ccm import calculate_ccm_from_ref
  8from .ccm_helper_functions import check_even_square, make_even_square
  9
 10# TODO: fix max_shift parameter
 11
 12def calculate_x_corr(im1: np.ndarray, im2: np.ndarray):
 13    ccm = calculate_ccm_from_ref(np.array([im2]).astype(np.float32), im1.astype(np.float32))[0]
 14
 15    return np.array(ccm)
 16
 17
 18def get_image_shift(im1: np.ndarray, im2: np.ndarray, box: int, max_shift: int = None):
 19    """Computes the shift from im1 to ima2"""
 20
 21    if (np.sum(im1) == 0) or (np.sum(im2) == 0):
 22        return 0, 0
 23
 24    # Compute image correlation
 25    x_corr = calculate_x_corr(im1, im2)
 26
 27    # crop XCorr based on max_shift
 28    w, h = im1.shape
 29    if max_shift > 0:
 30        x_border = int((w - max_shift) / 2)
 31        y_border = int((h - max_shift) / 2)
 32        if x_border > 0:
 33            x_corr = x_corr[x_border:-x_border, :]
 34        else:
 35            x_border = 0
 36        if y_border > 0:
 37            x_corr = x_corr[:, y_border:-y_border]
 38        else:
 39            y_border = 0
 40    else:
 41        x_border = y_border = 0
 42
 43    # A quarter of the fit ROI
 44    fit_box = int(box / 2)
 45
 46    # A coordinate grid for the fitting ROI
 47    x, y = np.mgrid[-fit_box : fit_box + 1, -fit_box : fit_box + 1]
 48
 49    # Find the brightest pixel and cut out the fit ROI
 50    x_max_xc, y_max_xc = np.unravel_index(x_corr.argmax(), x_corr.shape)
 51    fit_roi = x_corr[
 52        x_max_xc - fit_box : y_max_xc + fit_box + 1,
 53        y_max_xc - fit_box : y_max_xc + fit_box + 1,
 54    ]
 55
 56    dimensions = fit_roi.shape
 57
 58    if 0 in dimensions or dimensions[0] != dimensions[1]:
 59        xc, yc = 0, 0
 60    else:
 61        # The fit model
 62        def flat_2d_gaussian(a, xc, yc, s, b):
 63            A = a * np.exp(-0.5 * ((x - xc) ** 2 + (y - yc) ** 2) / s**2) + b
 64            return A.flatten()
 65
 66        gaussian2d = lmfit.Model(
 67            flat_2d_gaussian, name="2D Gaussian", independent_vars=[]
 68        )
 69
 70        # Set up initial parameters and fit
 71        params = lmfit.Parameters()
 72        params.add("a", value=fit_roi.max(), vary=True, min=0)
 73        params.add("xc", value=0, vary=True)
 74        params.add("yc", value=0, vary=True)
 75        params.add("s", value=1, vary=True, min=0)
 76        params.add("b", value=fit_roi.min(), vary=True, min=0)
 77        tmp = fit_roi.flatten()
 78        results = gaussian2d.fit(tmp, params)
 79
 80        # Get maximum coordinates and add offsets
 81        xc = results.best_values["xc"]
 82        yc = results.best_values["yc"]
 83        xc += y_border + x_max_xc
 84        yc += x_border + y_max_xc
 85
 86        xc -= np.floor(w / 2)
 87        yc -= np.floor(h / 2)
 88
 89    return -xc, -yc
 90
 91
 92def rcc(im_frames: np.ndarray, max_shift=None) -> tuple:
 93    # REF: https://github.com/yinawang28/RCC
 94    if not check_even_square(im_frames.astype(np.float32)):
 95        im_frames = np.array(make_even_square(im_frames.astype(np.float32)))
 96    n_frames = im_frames.shape[0]
 97    shifts_x = np.zeros((n_frames, n_frames))
 98    shifts_y = np.zeros((n_frames, n_frames))
 99    n_pairs = int(n_frames * (n_frames - 1) / 2)
100    flag = 0
101
102    with tqdm(
103        total=n_pairs, desc="Correlating image pairs", unit="pairs"
104    ) as progress_bar:
105        for i in range(n_frames - 1):
106            for j in range(i + 1, n_frames):
107                progress_bar.update()
108                shifts_x[i, j], shifts_y[i, j] = get_image_shift(
109                    im_frames[i], im_frames[j], 5, max_shift
110                )
111                flag += 1
112
113    return minimize_shifts(shifts_x, shifts_y)
114
115
116def minimize_shifts(shifts_x, shifts_y, shifts_z=None):
117    n_channels = shifts_x.shape[0]
118    n_pairs = int(n_channels * (n_channels - 1) / 2)
119    n_dims = 2 if shifts_z is None else 3
120    rij = np.zeros((n_pairs, n_dims))
121    A = np.zeros((n_pairs, n_channels - 1))
122    flag = 0
123    for i in range(n_channels - 1):
124        for j in range(i + 1, n_channels):
125            rij[flag, 0] = shifts_x[i, j]
126            rij[flag, 1] = shifts_y[i, j]
127            if n_dims == 3:
128                rij[flag, 2] = shifts_z[i, j]
129            A[flag, i:j] = 1
130            flag += 1
131    Dj = np.dot(np.linalg.pinv(A), rij)
132    shift_x = np.insert(np.cumsum(Dj[:, 0]), 0, 0)
133    shift_y = np.insert(np.cumsum(Dj[:, 1]), 0, 0)
134    if n_dims == 2:
135        return shift_x, shift_y
136    else:
137        shift_z = np.insert(np.cumsum(Dj[:, 2]), 0, 0)
138        return shift_x, shift_y, shift_z